import cv2
import matplotlib.pyplot as plt
import numpy as np
plt.rcParams['font.sans-serif'] = 'SimHei'


def plot_image_and_r2(image, x, y, k, b, r2):
    # 将 BGR 格式转换为 RGB 格式
    image = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2RGB)
    # 创建一个图形和两个子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5), gridspec_kw={'top': 0.85})
    # 设置窗口标题方式二
    fig.canvas.manager.window.title("建模结果")
    # 在第一个子图中显示图片
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title('Box')

    # 在第二个子图中显示拟合直线
    ax2.plot(x, y, 'o', label='Data')
    ax2.plot(x, k*x + b, label='Fitted Line')
    ax2.set_title('Fitted Line (k={:.5f}, b={:.5f} r2={:.5f})'.format(k, b,r2))
    ax2.legend()
    # 显示图形
    plt.tight_layout()
    plt.show()


def plot_image_and_r2_zzz(image, x, y,r2,theta):
    # 将 BGR 格式转换为 RGB 格式
    image = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2RGB)
    # 创建一个图形和两个子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5), gridspec_kw={'top': 0.85})
    # 设置窗口标题方式二
    fig.canvas.manager.window.title("建模结果")
    # 在第一个子图中显示图片
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title('Box')

    # 在第二个子图中显示拟合直线
    ax2.plot(x, y, 'o', label='Data')
    ax2.plot(x, x, label='Fitted Line')
    # 将每个数字转换为字符串，保留五位小数
    theta_str = "(k1={:.4f}, k2={:.4f}, b={:.4f})".format(*theta)
    ax2.set_title('Fitted Line (theta={}, r2={:.5f})'.format(theta_str,r2))
    # 添加轴标签
    ax2.set_xlabel('y_true')
    ax2.set_ylabel('y_pred')
    ax2.legend()
    # 显示图形
    plt.tight_layout()
    plt.show()

def plot_image_and_r2_3d(image, x, y, theta, r2):
    # 将 BGR 格式转换为 RGB 格式
    image = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2RGB)

    # 创建一个图形和两个子图
    fig = plt.figure(figsize=(12, 6))
    fig.canvas.manager.set_window_title("建模结果")

    # 在第一个子图中显示图片
    ax1 = fig.add_subplot(121)
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title('Box')

    # 在第二个子图中显示拟合平面
    ax2 = fig.add_subplot(122, projection='3d')
    x1,x2=x[:, 0], x[:, 1]
    ax2.scatter(x[:, 0], x[:, 1], y, color='b', label='Data')

    # 生成用于绘制回归面的网格数据
    x1_surf, x2_surf = np.meshgrid(np.linspace(x[:, 0].min(), x[:, 0].max(), 100),
                                   np.linspace(x[:, 1].min(), x[:, 1].max(), 100))
    z_surf = theta[0] + theta[1] * x1_surf + theta[2] * x2_surf

    # 绘制回归平面
    ax2.plot_surface(x1_surf, x2_surf, z_surf, color='r', alpha=0.5)
    ax2.set_title('Fitted Plane (R^2={:.5f})'.format(r2))
    ax2.set_xlabel('X1')
    ax2.set_ylabel('X2')
    ax2.set_zlabel('Y')

    # 显示图形
    plt.tight_layout()
    plt.show()

# 示例数据

if __name__ == '__main__':
    path = '20240531_113524\Point20_9.5.jpg'
    target = cv2.imread(path)
    # 调用方法绘制图像和 R2 图
    # 示例数据
    np.random.seed(0)
    X = np.random.rand(100, 2) * 10  # 生成100个随机的二维变量
    y = 2 * X[:, 0] + 3 * X[:, 1] + 5 + np.random.randn(100) * 2  # 生成对应的y值

    # 添加截距项
    X_b = np.c_[np.ones((100, 1)), X]  # 在 X 前面加一列1

    # 计算回归系数
    theta = np.linalg.inv(X_b.T @ X_b) @ X_b.T @ y

    # 计算 R^2 值
    y_pred = X_b @ theta
    ss_total = np.sum((y - np.mean(y)) ** 2)
    ss_res = np.sum((y - y_pred) ** 2)
    r2 = 1 - (ss_res / ss_total)

    # 示例图片
    image = np.zeros((100, 100, 3), dtype=np.uint8)
    plot_image_and_r2_3d(target, X, y, theta, 0.99)
